#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/11
@微信公众号: 邵奈一
"""

# 代码3-14
# 导入相关库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision

# 代码3-10
# 超参数设置
n_epochs = 3  # 定义训练周期数，这里设为3
batch_size_train = 64  # 定义训练批次大小，这里设为64
batch_size_test = 1000  # 定义测试批次大小，这里设为1000
learning_rate = 0.01  # 定义学习率，这里设为0.01
momentum = 0.5  # 定义动量，这里设为0.5
log_interval = 10  # 定义日志记录间隔，这里设为10
random_seed = 1  # 定义随机种子，这里设为1
torch.manual_seed(random_seed)  # 使用torch库的manual_seed函数设置随机种子，保证实验可复现

# 代码3-11
# 加载数据集
train_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('data3/',
                             train=True,    # 表示是训练集
                             download=False,     # 表示不进行下载操作（因为已经存在）
                             transform=torchvision.transforms.Compose([
                                 torchvision.transforms.ToTensor(),
                                 torchvision.transforms.Normalize((0.1307,), (0.3081,)) ])),
  batch_size=batch_size_train, shuffle=True)    # shuffle是指是否对取出的图片顺序进行打乱操作，以增加随机性从而提高模型泛化能力防止过拟合现象

test_loader = torch.utils.data.DataLoader(
  torchvision.datasets.MNIST('data3/',
                             train=False,
                             download=False,
                             transform=torchvision.transforms.Compose([
                                 torchvision.transforms.ToTensor(),
                                 torchvision.transforms.Normalize((0.1307,), (0.3081,)) ])),
  batch_size=batch_size_test, shuffle=True)

# 代码3-15
# 创建一个新类Net
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 定义第一个卷积层，输入通道数为1，输出通道数为10，卷积核大小为5
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        # 定义第二个卷积层，输入通道数为10，输出通道数为20，卷积核大小为5
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        # 定义一个Dropout层，用于随机丢弃部分神经元，防止过拟合
        self.conv2_drop = nn.Dropout2d()
        # 定义第一个全连接层，输入神经元数为320（4*4*20），输出神经元数为50
        self.fc1 = nn.Linear(320, 50)
        # 定义第二个全连接层，输入神经元数为50，输出神经元数为10
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        # 卷积层1的输出经过ReLU激活函数和最大池化层，池化核大小为2
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # 卷积层2的输出经过ReLU激活函数，同时随机丢弃部分神经元，再经过最大池化层，池化核大小为2
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        # 将卷积层2的输出展平
        x = x.view(-1, 320)
        # 全连接层1的输出经过ReLU激活函数
        x = F.relu(self.fc1(x))
        # 在训练过程中，随机丢弃部分神经元，防止过拟合（评估模式下self.training=false）
        x = F.dropout(x, training=self.training)
        # 全连接层2的输出
        x = self.fc2(x)
        # 对全连接层的输出进行log_softmax处理，得到最后的输出结果
        return F.log_softmax(x)


# 代码3-16
# 初始化网络和优化器
network = Net()
# 实例化一个优化器对象，使用随机梯度下降（SGD）算法
# 参数network.parameters()：网络的参数
# lr=learning_rate：学习率，控制模型学习的速度
# momentum=momentum：动量，控制模型学习的稳定性
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)

# 代码3-17
# 参数设置
# 初始化训练损失列表
train_losses = []
# 初始化训练计数器列表
train_counter = []
# 初始化测试损失列表
test_losses = []
# 初始化测试计数器列表，其元素为从0到n_epochs的每个元素乘以训练数据集的长度
test_counter = [i * len(train_loader.dataset) for i in range(0, n_epochs + 1)]


# 代码3-18
# 一次测试循环
def train(epoch):
    # 将网络设置为训练模式
    network.train()
    # 遍历训练数据加载器
    for batch_idx, (data, target) in enumerate(train_loader):
        # 每个批次开始时，清零梯度
        optimizer.zero_grad()
        # 前向传播，获取模型的输出
        output = network(data)
        # 计算损失函数，这里使用负对数似然损失（NLL Loss）
        loss = F.nll_loss(output, target)
        # 反向传播，计算梯度
        loss.backward()
        # 更新模型参数
        optimizer.step()
        # 每隔 log_interval 个批次打印一次训练信息（batch_idx为批次，0，10，20，30...的时候会打印，data是一个批次的数据，每个批次是64）
        if batch_idx % log_interval == 0:
            # {{:.6f}} 中的冒号用于指定格式说明符，表示浮点数格式化。
            # 100. * batch_idx 中的句号表示浮点数，确保运算结果为浮点数类型。
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))
            # 记录训练损失
            train_losses.append(loss.item())
            # 记录训练步数
            train_counter.append((batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
            # 保存模型和优化器的状态字典
            torch.save(network.state_dict(), 'tmp/3-model.pth')
            torch.save(optimizer.state_dict(), 'tmp/3-optimizer.pth')
            # 输出结果：
            # Train Epoch: 3 [58240/60000 (97%)]	Loss: 0.105537
            # Train Epoch: 3 [58880/60000 (98%)]	Loss: 0.276362
            # Train Epoch: 3 [59520/60000 (99%)]	Loss: 0.242384

# train(1)


# 代码3-19
# 测试循环
def test():
    network.eval()
    # 初始化测试损失和正确预测的数量
    test_loss = 0
    correct = 0
    with torch.no_grad(): # 上下文管理器禁用梯度计算，以节省内存和加速计算
        for data, target in test_loader:
            output = network(data)
            # 计算测试损失，注意 size_average=False 表示不平均损失
            test_loss += F.nll_loss(output, target, size_average=False).item()
            # 获取预测的类别
            pred = output.data.max(1, keepdim=True)[1]
            # 计算正确预测的数量
            correct += pred.eq(target.data.view_as(pred)).sum()
    # 计算平均测试损失
    test_loss /= len(test_loader.dataset)
    # 记录测试损失
    test_losses.append(test_loss)
    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    # 输出结果：
    # Test set: Avg. loss: 0.1021, Accuracy: 9688/10000 (97%)

# test()

# 代码3-20
# 模型评估
test()
for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

# 代码3-21
# 性能评估
import matplotlib.pyplot as plt

# 设置字体
plt.rcParams['font.family'] = 'Heiti TC'

fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.plot(test_counter, test_losses, color='red', linestyle=':')
plt.legend(['训练损失', '测试损失'], loc='upper right')
plt.xlabel('训练次数')
plt.ylabel('损失')
plt.savefig('tmp/cnn-number.png', dpi=1080)
plt.show()

# 代码3-22
# 预测结果
import matplotlib.pyplot as plt
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
with torch.no_grad():
    output = network(example_data)
fig = plt.figure()
for i in range(6):
    plt.subplot(2, 3, i + 1)
    plt.tight_layout()
    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
    plt.title('预测结果: {}'.format(
        output.data.max(1, keepdim=True)[1][i].item())) # keepdim=True 表示在计算最大值后，输出张量的维度将保持与输入张量相同
    plt.xticks([])
    plt.yticks([])
    plt.savefig('tmp/cnn-number1.png', dpi=1080)
plt.show()